package com.zhuoli.service.springboot.local.current.limit.aop;

import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import com.zhuoli.service.springboot.local.current.limit.annotation.Limit;
import com.zhuoli.service.springboot.local.current.limit.common.enums.LimitType;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Date;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @Author: zhuoli
 * @Date: 2018/9/19 10:23
 * @Description:
 */
@Aspect
@Configuration
@Slf4j
@RequiredArgsConstructor
public class LimitInterceptor {

    private static final Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    @Around("execution(public * *(..)) && @annotation(com.zhuoli.service.springboot.local.current.limit.annotation.Limit)")
    public Object interceptor(ProceedingJoinPoint pjp) {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();
        String name = limitAnnotation.name();
        String key;
        int limitCount = limitAnnotation.count();
        int expireTime = limitAnnotation.expireTime();
        switch (limitType) {
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = method.getName().toUpperCase();
        }

        String keyRes = Joiner.on("").join(limitAnnotation.prefix(), key);
        try {
            if (Objects.isNull(rateLimiterMap.get(keyRes))) {
                rateLimiterMap.put(keyRes, RateLimiter.create(limitCount));
            }
            RateLimiter rateLimiter = rateLimiterMap.get(keyRes);

            /*非阻塞式，预计1S内获取不到服务，则不继续等待，直接返回null*/
            if (!rateLimiter.tryAcquire(expireTime, TimeUnit.SECONDS)) {
                log.info("Waiting time will be too long, please try later, now is {}", new Date());
                return null;
                //为了测试时方便看控制台日志，把下面代码注释，实际使用将return null注释，将下面代码释放
                //throw new RuntimeException("can not get service in expireTime,please try later");
            }
            log.info(keyRes + " : 成功 : " + new Date());

            //阻塞式，请求会被阻塞，直到获取令牌，采取执行pjp.proceed()
            //log.info(keyRes + " : 成功 : " + "等待时间 : " + rateLimiter.acquire(1));

            return pjp.proceed();
        } catch (Throwable e) {
            if (e instanceof RuntimeException) {
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
    }

    private static final String UNKNOWN = "unknown";

    public String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}
